#!/usr/bin/python3
import json
import signal
import sys
import ipaddress
import random
import os
from ipaddress import ip_address
from time import sleep

from scapy.contrib.openflow import ofp_stats_reply_cls

IPV4 = ipaddress.IPv4Address._ALL_ONES
IPV6 = ipaddress.IPv6Address._ALL_ONES
ULP = ['tcp', 'udp']
SPI = range(10, 2**16 + 10)
SPI_INDEX = 0


def quit(sigal_num, frame):
    print()
    print("Ctrl + C----------Stop")
    sys.exit()

def NumberFilter(s_origin: str) -> str:
    s_res = ''
    for c in s_origin:
        if c.isnumeric():
            s_res = s_res + c
    return s_res


def random_ipv4():
    return ipaddress.IPv4Address._string_from_ip_int(
        random.randint(0, IPV4)
    )


def random_ipv6():
    return ipaddress.IPv6Address._string_from_ip_int(
        random.randint(0, IPV6)
    )


def GetRandomIP(ip_type: int = 4) -> str:
    if ip_type == 4:
        return random_ipv4()
    elif ip_type == 6:
        return random_ipv6()
    else:
        return ''

def SetSpiRandom(num : int):
    global SPI
    SPI = random.sample(SPI, num * 2)

def GetSpi()->int:
    global SPI, SPI_INDEX
    res = SPI[SPI_INDEX]
    SPI_INDEX = SPI_INDEX + 1
    return res
def GetAnEntryPair() -> tuple:
    bd_ipv4 = GetRandomIP(4)
    bd_ipv6 = GetRandomIP(6)
    dd_ipv4 = GetRandomIP(4)
    dd_ipv6 = GetRandomIP(6)
    sport = random.randint(10000, 50000)
    dport = random.randint(10000, 50000)
    ulp = ULP[random.randint(0, 1)]
    spi_tx = GetSpi()
    spi_rx = GetSpi()
    return bd_ipv4, dd_ipv4, spi_tx, spi_rx, sport, dport, bd_ipv6, dd_ipv6, ulp

def AddSpSa(entry : list, eth, io_type, offload : bool = False):
    if io_type == '4-6':
        inner_sip = entry[6]
        inner_dip = entry[7]
        outer_sip = entry[0]
        outer_dip = entry[1]
        ip_mask = '64'
    elif io_type == '6-4':
        inner_sip = entry[0]
        inner_dip = entry[1]
        outer_sip = entry[6]
        outer_dip = entry[7]
        ip_mask = '24'
    if offload:
        offload_str = "offload packet dev {}".format(eth)
        dir_in_str = "dir in"
        dir_out_str = "dir out"
    else:
        offload_str = ""
        dir_in_str = ""
        dir_out_str = ""
    spi_tx = entry[2]
    spi_rx = entry[3]
    sport = entry[4]
    dport = entry[5]
    ulp = entry[8]

    cmd1 = "ip xfrm state add src {} dst {} proto esp spi {} mode tunnel reqid {} aead 'rfc4106(gcm(aes))' 0x444342413433323124232221141312112f8d4454 64 " \
        "sel src {}/{} dst {}/{} sport {} dport {} proto {} replay-window 32 {} {}"
    cmd2 = "ip xfrm state add src {} dst {} proto esp spi {} mode tunnel reqid {} aead 'rfc4106(gcm(aes))' 0x44434241343332312423222114131211117e050a 64 " \
           "sel src {}/{} dst {}/{} sport {} dport {} proto {} replay-window 32 {} {}"
    cmd3 = "ip xfrm policy add dir in src {}/{} dst {}/{} sport {} dport {} proto {} tmpl proto esp src {} dst {} spi {} mode tunnel reqid {} {}"
    cmd4 = "ip xfrm policy add dir out src {}/{} dst {}/{} sport {} dport {} proto {} tmpl proto esp src {} dst {} spi {} mode tunnel reqid {} {}"

    cmd1 = cmd1.format(outer_dip, outer_sip, spi_rx, spi_rx, inner_dip, ip_mask, inner_sip, ip_mask, dport, sport, ulp, offload_str, dir_in_str)
    cmd2 = cmd2.format(outer_sip, outer_dip, spi_tx, spi_tx, inner_sip, ip_mask, inner_dip, ip_mask, sport, dport, ulp, offload_str, dir_out_str)
    cmd3 = cmd3.format(inner_dip, ip_mask, inner_sip, ip_mask, dport, sport, ulp, outer_dip, outer_sip, spi_rx, spi_rx, offload_str)
    cmd4 = cmd4.format(inner_sip, ip_mask, inner_dip, ip_mask, sport, dport, ulp, outer_sip, outer_dip, spi_tx, spi_tx, offload_str)

    os.system(cmd1)
    os.system(cmd2)
    os.system(cmd3)
    os.system(cmd4)
    #
    # print()
    # print(cmd1)
    # print(cmd2)
    # print(cmd3)
    # print(cmd4)

if __name__ == '__main__':

    signal.signal(signal.SIGINT, quit)
    signal.signal(signal.SIGTERM, quit)

    argv = len(sys.argv)
    args = sys.argv
    dict_path = 'entry.dict'

    # test
    # argv = 5
    # args = ['', '-s', 'wlp2s0', '4-6', '']

    if argv <= 2:
        print("Must have one parameter given.")
        exit(0)

    do = args[1]

    if do == '-g':
        if argv != 3:
            print("Must have two parameters given.")
            exit(0)

        value = 0
        num = args[2]

        if num.endswith('k') or num.endswith('K'):
            value = int(NumberFilter(num)) * 1024
        else:
            value = int(NumberFilter(num))

        entries_to_write = []
        SetSpiRandom(value)
        num = 0
        while True:
            entry = GetAnEntryPair()
            if entry in entries_to_write:
                continue
            else:
                entries_to_write.append(entry)
                num = num + 1
                if num >= value:
                    break
        with open(dict_path, 'w') as f:
            f.write(json.dumps(entries_to_write))
    elif do == '-s':
        if argv < 4:
            print("Must have three/four parameters given.")
            exit(0)
        if not os.path.exists(dict_path):
            print("No dict found.Please exec: python3 xfrm_test.py -g [num] before this.")
            exit(0)
        eth = args[2]
        io_type = args[3]
        offload = False
        if argv == 5 :
            if args[4] == 'offload':
                offload = True

        with open(dict_path, 'r') as f:
            raw_data = f.read()
        entries = json.loads(raw_data)

        # entries[0] = (sip4, dip4, spi_tx, spi_rx, sport, dport, sip6, dip6, ulp)
        # clean
        os.system("ip x p f")
        os.system("ip x s f")

        length = len(entries)
        for i in range(1, length + 1):
            AddSpSa(entries[i - 1], eth, io_type, offload)
            print("\rExecute NO.{}\t\tProcess:{}%".format(i, round(float(i) / length * 100.0, 3)), end="")
            sleep(0.1)

        print("\nDone!")

    elif do == '-h':
        print("xfrm_test v0.1")
        print("You can do like these:")
        print("\t\t\t\t-g [num][k]: generate num[k] SP entries dict")
        print("\t\t\t\t-s [device] [4-6/6-4] [null/offload]: set some(the num you have given) sp, and num * 2 sa and config offload")
        print("\r\n")
    else:
        print("Invalid input.")
        exit(0)
